import tensorflow as tf
from sklearn import datasets
import numpy as np
import matplotlib.pyplot as plt
# 数据集注入
x_data=datasets.load_iris().data  # 获取鸢尾花的所有输入特征
y_data=datasets.load_iris().target  # 获取对应的鸢尾花标签集
count=len(x_data)

# 数据集乱序,输入特征和标签使用相同的随机数种子保证两者有一一对应
np.random.seed(114)
np.random.shuffle(x_data)
np.random.seed(114)
np.random.shuffle(y_data)
tf.random.set_seed(116)  # 这一句是干嘛？？？？

# 划分训练、测试数据集
x_train=x_data[:(int)(count*0.8)]
x_test=x_data[-(int)(count*0.2):]
y_train=y_data[:(int)(count*0.8)]
y_test=y_data[-(int)(count*0.2):]

# 转换x的数据类型，否则后面矩阵相乘时会因数据类型不一致报错
x_train = tf.cast(x_train, tf.float32)
x_test = tf.cast(x_test, tf.float32)

# 结合成（特征，标签）数据集,并且32数据为一个个小组
train_ds=tf.data.Dataset.from_tensor_slices((x_train,y_train)).batch(32)
test_ds=tf.data.Dataset.from_tensor_slices((x_test,y_test)).batch(32)

print(train_ds)
# 定义神经网络的所有可训练参数,此处只有一层神经网络，所以4是输入层的特征维度4（花萼长度', '花萼宽度', '花瓣长度', '花瓣宽度'），3输出层，三种鸢尾花的匹配率
w1 = tf.Variable(tf.random.truncated_normal([4,3],stddev=0.1,seed=1))
b1 = tf.Variable(tf.random.truncated_normal([3],stddev=0.1,seed=1))

lr = 0.1  # 学习率为0.1
train_loss_results = []  # 将每轮的loss记录在此列表中，为后续画loss曲线提供数据
test_acc = []  # 将每轮的acc记录在此列表中，为后续画acc曲线提供数据
epoch = 500  # 循环500轮
loss_all = 0  # 每轮分4个step，loss_all记录四个step生成的4个loss的和

# 嵌套迭代,输出当前的loss
for epoch in range(epoch):  #epoch 级别迭代
    for step,(x_train,y_train) in enumerate(train_ds):  #batch 级别的迭代
        with tf.GradientTape() as tape: #记录梯度信息
            y=tf.matmul(x_train,w1)+b1 # 神经网络乘加运算
            y=tf.nn.softmax(y) # 使得输出y符合的概率分布
            y_=tf.one_hot(y_train,depth=3) # 将标签转换成3维的独热编码
            loss=tf.reduce_mean(tf.square(y-y_)) # 采用均方误差损失函数mse = mean(sum(y-out)^2)
            loss_all+=loss.numpy()
        # 计算loss 函数对于各个参数的梯度
        grads=tape.gradient(loss,[w1,b1])

        # 实现梯度更新 w1 = w1 - lr * w1_grad    b = b - lr * b_grad
        w1.assign_sub(lr*grads[0])
        b1.assign_sub(lr*grads[1])

    # 每个epoch，打印loss信息
    print("Epoch {}, loss: {}".format(epoch, loss_all / 4))
    train_loss_results.append(loss_all / 4)  # 将4个step的loss求平均记录在此变量中
    loss_all = 0  # loss_all归零，为记录下一个epoch的loss做准备

    # 测试数据部分
    total_correct=0  # 表示模型预测正确的个数
    total_num=0  # 表示总的测试个数
    # 遍历测试数据集
    for x_test,y_test in test_ds:
        # 使用上面训练的参数进行预测
        y=tf.matmul(x_test,w1)+b1
        y=tf.nn.softmax(y) # 返回一个数组：包含四种鸢尾花的匹配概率
        predict=tf.argmax(y,axis=1) # 取概率最高的即为本次的预测结果
        predict=tf.cast(predict,dtype=y_test.dtype) #将预测结果转换成独热编码，方便比对是否预测成功
        # 将对比结果由布尔型转变成0和1，方便正确个数的计数
        correct=tf.cast(tf.equal(predict,y_test),dtype=tf.int32)
        #将每个batch的正确结果统计起来
        correct=tf.reduce_sum(correct)
        # 将所有batch中的correct数加起来
        total_correct+=correct
        # total_num为测试的总样本数，也就是x_test的行数，shape[0]返回变量的行数
        total_num += x_test.shape[0]

    # 总的准确率等于total_correct/total_number
    acc = total_correct / total_num
    test_acc.append(acc)
    print("Test_acc:", acc)
    print("--------------------------")

# 评估模型，建立相关的图表
# 绘制 loss 曲线
plt.title('Loss Function Curve')  # 图片标题
plt.xlabel('Epoch')  # x轴变量名称
plt.ylabel('Loss')  # y轴变量名称
plt.plot(train_loss_results, label="$Loss$")  # 逐点画出trian_loss_results值并连线，连线图标是Loss
plt.legend()  # 画出曲线图标
plt.show()  # 画出图像

# 绘制 Accuracy 曲线
plt.title('Acc Curve')  # 图片标题
plt.xlabel('Epoch')  # x轴变量名称
plt.ylabel('Acc')  # y轴变量名称
plt.plot(test_acc, label="$Accuracy$")  # 逐点画出test_acc值并连线，连线图标是Accuracy
plt.legend()
plt.show()

